import numpy as np
import quimb.tensor as qtn

"""the setting of TensorNetwork QPT"""
N_qubit: int = 10
N_set: int = 10
N_shot: int = 1000

from splm_coeff import even_CNOT_10qubit_lam1s, even_CNOT_10qubit_lam2s

real_lam1_mat = even_CNOT_10qubit_lam1s
real_lam2_mat = even_CNOT_10qubit_lam2s

# filename of the measurement data of simulated sampling
import json
import pathlib

filename = f"QPT_dataset1 N_qubit={N_qubit}, N_shot={N_shot}.json"
file = pathlib.Path(filename)
if file.exists():
	with open(file, "r") as f:
		metadata = json.load(f)["metadata"]
	lam1_mat, lam2_mat = metadata["lam1_mat"], metadata["lam2_mat"]
	conds = metadata["N_qubit"] == N_qubit and metadata["N_shot"] == N_shot \
	        and metadata["rhoin_set"] == "xyz proj" and metadata["meas_set"] == "xyz" \
	        and np.isclose(lam1_mat, real_lam1_mat).all() and np.isclose(lam2_mat, real_lam2_mat).all()

	if not conds: raise FileExistsError

"""
电路某层的酉算符的张量网络表示 这里采用相邻比特的CNOT门作为一层电路
Sparse Pauli Lindblad Model的PTM表示
"""
from QPT_utils import even_CNOT_mpo
from TNQPT.noise_model import PTM_SPLM_MPO

ideal_mpo = even_CNOT_mpo(N_qubit,
                          upper_ind_id="a{}",
                          lower_ind_id="b{}",
                          site_tag_id="I{}")

onchip_splm_mpo = PTM_SPLM_MPO.from_adj_param_mat(
	real_lam1_mat, real_lam2_mat,
	upper_ind_id="k{}", lower_ind_id="a{}", site_tag_id="I{}"
)

noisy_mpo = qtn.tensor_network_1d_compress(ideal_mpo | onchip_splm_mpo,
                                           method='direct', max_bond=400, cutoff=1e-10, )
"""
准备N_set×N_shot个输入态和对应的测量观测数据
对于每一组N_set，给定输入密度矩阵和对应的pauli测量, 对于每一个qubit位置, 随机分配一个{X, Y, Z}pauli测量算符
prepare MPO representation of input density operator and corresponding measurement settings
"""
from TNQPT.states import ptm_pauli_proj
from TNQPT.data_gen import QPT_data_gen


Xp_o1 = ptm_pauli_proj('x+').reshape(-1)
Xm_o1 = ptm_pauli_proj('x-').reshape(-1)
Yp_o1 = ptm_pauli_proj('y+').reshape(-1)
Ym_o1 = ptm_pauli_proj('y-').reshape(-1)
Zp_o1 = ptm_pauli_proj('z+').reshape(-1)
Zm_o1 = ptm_pauli_proj('z-').reshape(-1)

rhoin_map = {
	'x+': Xp_o1, 'x-': Xm_o1,
	'y+': Yp_o1, 'y-': Ym_o1,
	'z+': Zp_o1, 'z-': Zm_o1,
}
meas_map = {
	'x': {'x+': Xp_o1, 'x-': Xm_o1},
	'y': {'y+': Yp_o1, 'y-': Ym_o1},
	'z': {'z+': Zp_o1, 'z-': Zm_o1}
}

rhoin_settings = np.random.choice(list(rhoin_map.keys()), size=(N_set, N_qubit))
meas_settings = np.random.choice(list(meas_map.keys()), size=(N_set, N_qubit))


datas = [
	list(QPT_data_gen(N_shot,
	                  noisy_mpo,
	                  rhoin_map,
	                  meas_map,
	                  rhoin_setting,
	                  meas_setting,
	                  desc=f"Sampling {i + 1}/{N_set}"))
	for i, (rhoin_setting, meas_setting) in enumerate(zip(rhoin_settings, meas_settings))
]

"""saving the measurement data"""

if file.exists():
	with open(file, "r") as f:
		data_dict = json.load(f)
	data_dict["metadata"]["N_set"] += N_set
	data_dict["samples"]["rhoin"].extend(rhoin_settings.tolist())
	data_dict["samples"]["meas"].extend(meas_settings.tolist())
	data_dict["samples"]["datas"].extend(datas)
else:
	data_dict = {
		"description": "the dataset of quantum tomograph of even CNOT gate",
		"metadata": {
			"N_qubit": N_qubit,
			"N_set": N_set,
			"N_shot": N_shot,
			"rhoin_set": "xyz proj",
			"meas_set": "xyz",

			"lam1_mat": real_lam1_mat.tolist(),
			"lam2_mat": real_lam2_mat.tolist(),
		},
		"samples": {
			'rhoin': rhoin_settings.tolist(),
			'meas': meas_settings.tolist(),
			'datas': datas
		}
	}

with open(filename, 'w') as f:
	json.dump(data_dict, f)
